#pragma once

#include "vkinfer/runtime/blob.h"
#include <vulkan/vulkan.h>
#include <memory>
#include "vkinfer/runtime/vk/vktransfer.h"

namespace vkinfer
{
    struct VkBlob : public Blob
    {
        VkDevice device = VK_NULL_HANDLE;
        VkPhysicalDevice gpu = VK_NULL_HANDLE;
        std::shared_ptr<VkTransfer> transfer;

        VkBuffer buffer = VK_NULL_HANDLE;
        VkDeviceMemory memory = VK_NULL_HANDLE;
        VkBuffer stagingBuffer = VK_NULL_HANDLE;
        VkDeviceMemory stagingMemory = VK_NULL_HANDLE;
        VkDeviceSize size = 0;
        BlobTransferDirection direction = BlobTransferDirection::BLOB_TRANSFER_NONE;

        VkBlob(VkDevice device, VkPhysicalDevice gpu, const std::shared_ptr<VkTransfer>& transfer,
            BlobTransferDirection direction);

        VkBlob(VkDevice device, VkPhysicalDevice gpu, const std::shared_ptr<VkTransfer>& transfer, 
            const std::vector<uint32_t>& shape, BlobTransferDirection direction);
        ~VkBlob();
        virtual void copy_from(const std::shared_ptr<Tensor<float>>& tensor);
        virtual void copy_from(const std::vector<float>& data) override;
        virtual void copy_from(const float* data, const std::vector<uint32_t>& shape) override;
        virtual void copy_to(std::vector<float>& data) override;
        virtual void resize(const std::vector<uint32_t>& shape) override;

    private:
        void destory();
        void create(const std::vector<uint32_t>& shape);
    };
}
